#!/usr/bin/env python3
import os
import argparse
import pysam
from .remove_no_chimeric import IvalTree
from .utils import is_overlap, reverse_seq


class Transcript:
    def __init__(self, transcript_id, transcript_name, transcript_type, gene_id,
                 gene_name, gene_type, strand, chrom, tx_start, tx_end, cds_start,
                 cds_end, exon_sizes, exon_starts):
        self.transcript_id = transcript_id
        self.transcript_name = transcript_name
        self.transcript_type = transcript_type
        self.gene_id = gene_id
        self.gene_name = gene_name
        self.gene_type = gene_type
        self.chrom = chrom
        self.strand = strand
        self.tx_start = int(tx_start)
        self.tx_end = int(tx_end)
        self.cds_start = int(cds_start)
        self.cds_end = int(cds_end)
        self.exon_sizes = exon_sizes
        self.exon_starts = exon_starts

    def __str__(self):
        return f"{self.transcript_id} {self.transcript_name} {self.transcript_type} {self.strand}"

    def __repr__(self):
        return f"Transcript({self.transcript_id}"

    @property
    def exons(self):
        exon_range = []
        exon_starts = list(map(int, self.exon_starts.rstrip(',').split(',')))
        exon_starts = list(map((lambda x: x + self.tx_start), exon_starts))
        exon_ends = list(map(int, self.exon_sizes.rstrip(',').split(',')))
        exon_ends = list(map((lambda x, y: x + y), exon_starts, exon_ends))
        for st, end in zip(exon_starts, exon_ends):
            exon_range.append([st, end])
        return exon_range

    @property
    def introns(self):
        intron_range = []
        exon_range = self.exons
        exon_starts = [i[0] for i in exon_range]
        exon_ends = [i[1] for i in exon_range]
        Intron_starts = [self.tx_start] + exon_ends
        exon_starts.append(self.tx_end)
        Intron_ends = exon_starts
        for st, end in zip(Intron_starts, Intron_ends):
            if st == end:
                continue
            intron_range.append([st, end])

        return intron_range

    def classify_exons(self):
        utr5_exons = []
        utr3_exons = []
        cds_exons = []
        if not self.transcript_type == 'protein_coding':
            return None
        assert self.cds_end > self.cds_start
        for exon in self.exons:
            if exon[1] < self.cds_start:
                utr5_exons.append(exon)
            elif exon[0] > self.cds_end:
                utr3_exons.append(exon)
            elif self.cds_start <= exon[0] and exon[1] <= self.cds_end:
                cds_exons.append(exon)
            elif exon[0] < self.cds_start < exon[1] and exon[1] <= self.cds_end:
                utr5_exons.append([exon[0], self.cds_start])
                cds_exons.append([self.cds_start, exon[1]])
            elif self.cds_start <= exon[0] and exon[0] < self.cds_end < exon[1]:
                cds_exons.append([exon[0], self.cds_end])
                utr3_exons.append([self.cds_end, exon[1]])
            elif exon[0] < self.cds_start and self.cds_end < exon[1]:
                utr5_exons.append([exon[0], self.cds_start])
                cds_exons.append([self.cds_start, self.cds_end])
                utr3_exons.append([self.cds_end, exon[1]])

        if self.strand == '-':
            utr5_exons, utr3_exons = utr3_exons, utr5_exons

        return utr5_exons, cds_exons, utr3_exons

    def anno_read(self, blocks):
        desc = ''
        blocks_start = blocks[0][0]
        blocks_end = blocks[-1][1]
        if blocks_start < self.tx_start or blocks_end > self.tx_end:
            assert is_overlap((blocks_start, blocks_end), (self.tx_start, self.tx_end))
            desc += "Outside_region,"

        for intron in self.introns:
            flag = False
            for interval in blocks:
                if is_overlap(intron, interval):
                    if self.transcript_type == 'protein_coding':
                        desc += "mRNA-intron,"
                    else:
                        desc += "intron,"
                    flag = True
                    break
            if flag:
                break
        if self.transcript_type == 'protein_coding':
            utr5_exons, cds_exons, utr3_exons = self.classify_exons()
            for exon in utr5_exons:
                flag = False
                for interval in blocks:
                    if is_overlap(exon, interval):
                        desc += '5utr_exon,'
                        flag = True
                        break
                if flag:
                    break
            for exon in utr3_exons:
                flag = False
                for interval in blocks:
                    if is_overlap(exon, interval):
                        desc += '3utr_exon,'
                        flag = True
                        break
                if flag:
                    break
            for exon in cds_exons:
                flag = False
                for interval in blocks:
                    if is_overlap(exon, interval):
                        desc += 'cds_exon,'
                        flag = True
                        break
                if flag:
                    break
        else:
            for exon in self.exons:
                flag = False
                for interval in blocks:
                    if is_overlap(exon, interval):
                        desc += 'exon,'
                        flag = True
                        break
                if flag:
                    break
        desc = desc.rstrip(',')
        desc_list = desc.split(",")
        desc_set = set(desc_list)
        assert len(desc_list) == len(desc_set)
        return desc

    @property
    def utr5_exons(self):
        if not self.transcript_type == 'protein_coding':
            return None
        utr5_exons, _, _ = self.classify_exons()

        return utr5_exons

    @property
    def utr3_exons(self):
        if not self.transcript_type == 'protein_coding':
            return None
        _, _, utr3_exons = self.classify_exons()

        return utr3_exons

    @property
    def cds_exons(self):
        if not self.transcript_type == 'protein_coding':
            return None
        _, cds_exons, _ = self.classify_exons()

        return cds_exons
    
    def get_sequence(self, genomefetcher):
        return genomefetcher.fetch(self.chrom, self.exons, self.strand)
    
    def get_sequence_with_intron(self, genomefetcher):
        return genomefetcher.fetch(self.chrom, [(self.tx_start, self.tx_end)], self.strand)


def main(trans_info, bed12, genome, pirna, bamfile, fasta, rna_anno):
    bam = pysam.AlignmentFile(bamfile, "rb")
    with open(trans_info) as t, open(bed12) as b, open(fasta, "w") as f, open(rna_anno, "w") as o:
        next(t)
        header = "read_name\tchrom\tstart\tend\tblocks\tstrand\tposition\ttranscript_id\ttranscript_type\tgene_name\tgene_type\n"
        o.write(header)
        transcript_tree = IvalTree()

        for info, bed in zip(t, b):
            info_fileds = info.rstrip().split("\t")
            bed_fileds = bed.rstrip().split("\t")
            assert info_fileds[0] == bed_fileds[3]
            transcript = Transcript(
                *info_fileds, bed_fileds[0], bed_fileds[1], bed_fileds[2], bed_fileds[6], bed_fileds[7], bed_fileds[10], bed_fileds[11])
            transcript_tree.add(transcript, transcript.chrom, transcript.tx_start, transcript.tx_end)

        for aligned_read in bam:
            if aligned_read.is_unmapped:
                continue
            if aligned_read.is_qcfail:
                continue
            if aligned_read.is_secondary:
                continue
            read2strand = '+' if aligned_read.is_reverse else '-'
            read_chrom = aligned_read.reference_name
            reference_start = aligned_read.reference_start
            reference_end = aligned_read.reference_end
            blocks = aligned_read.get_blocks()
            assert blocks[0][0] == reference_start
            assert blocks[-1][1] == reference_end
            read_name = aligned_read.query_name

            alignedTranscripts = transcript_tree.find(read_chrom, reference_start, reference_end)
            anno_count = 0
            if alignedTranscripts:
                for aligned_transcript in alignedTranscripts:
                    if not read2strand == aligned_transcript.strand:
                        continue
                    position = aligned_transcript.anno_read(blocks)
                    anno = (f"{read_name}\t{aligned_transcript.chrom}\t{reference_start}\t{reference_end}\t{blocks}\t{aligned_transcript.strand}\t"
                            f"{position}\t{aligned_transcript.transcript_id}\t{aligned_transcript.transcript_type}\t{aligned_transcript.gene_name}\t{aligned_transcript.gene_type}\n")
                    o.write(anno)
                    anno_count += 1
            else:
                anno = (f"{read_name}\t{read_chrom}\t{reference_start}\t{reference_end}\t{blocks}\t{read2strand}\t"
                        f"Unknown\tUnknown\tUnknown\tUnknown\tUnknown\n")
                o.write(anno)
                anno_count += 1

            if anno_count:
                sequence = aligned_read.query_alignment_sequence if read2strand == '+' else reverse_seq(aligned_read.query_alignment_sequence)
                f.write(f">{read_name}\n{sequence}\n")


def run():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-b", dest="bam",
                        type=str, required=True, help="bamfile")
    parser.add_argument("-g", dest="bed",
                        type=str, required=True, help="bed12")
    parser.add_argument("-a", dest="genome",
                        type=str, default=None, help="genome fasta")
    parser.add_argument("-i", dest="info",
                        type=str, required=True, help="transcript info")
    parser.add_argument("-f", dest="fasta", default=os.devnull,
                        type=str, help="mapped reads")
    parser.add_argument("-o", dest="rna_anno", default=os.devnull,
                        type=str, help="rna anno")
    args = parser.parse_args()

    main(args.info, args.bed, args.genome, args.bam, args.fasta, args.rna_anno)


if __name__ == "__main__":
    run()
